# -*- coding: utf-8 -*-

from pymongo import UpdateOne
from pe.util.DataBase import DB_CONN
from pe.util.StockUtil import getTradingDates
import tushare as ts
from datetime import datetime, timedelta
import traceback

"""
从tushare获取股票基础数据，保存到本地的MongoDB数据库中
"""

class BasicCrawler:

    def __init__(self):
        self.db = DB_CONN['basic']

    """
    抓取指定时间范围内的股票基础信息
    :param startDate: 开始日期
    :param endDate: 结束日期
    """
    def crawlBasic(self, startDate=None, endDate=None):

        if startDate is None:
            startDate = (datetime.now() - timedelta(days=1)).strftime('%Y-%m-%d')

        if endDate is None:
            endDate = (datetime.now() - timedelta(days=1)).strftime('%Y-%m-%d')

        allDates = getTradingDates(startDate, endDate)

        for date in allDates:
            try:
                self.crawlBasicAtDate(date)
            except:
                print('抓取股票基本信息时出错， 日期%s' % date, flush=True)

    """
    从Tushare抓取指定日期的股票基本信息
    :param date: 日期
    """
    def crawlBasicAtDate(self, date):
        # 默认推送上一个交易日的数据
        dfBasics = ts.get_stock_basics(date)

        # 如果当日没有基础信息，则不做操作
        if dfBasics is None:
            return

        updateRequests = []
        codes = set(dfBasics.index)

        for code in codes:
            doc = dict(dfBasics.loc[code])

            try:
                # 将20190101转换为2019-01-01的形式
                timeToMarket = datetime.strptime(str(doc['timeToMarket']), '%Y%m%d').strftime('%Y-%m-%d')

                # 解决流通股本和总股本单位不一致的情况
                totals = float(doc['totals'])

                # 这里假设最大规模的股本不超过5000亿，股本规模的最大工商银行是3564亿
                if totals > 5000:
                    totals *= 1E4
                else:
                    totals *= 1E8

                outstanding = float(doc['outstanding'])

                if outstanding > 5000:
                    outstanding *= 1E4
                else:
                    outstanding *= 1E8

                # 保存时增加date字段,因为每天都会有一条数据
                doc.update({
                    'code': code,
                    'date': date,
                    'timeToMarket': timeToMarket,
                    'outstanding': outstanding,
                    'totals': totals
                })

                updateRequests.append(
                    UpdateOne(
                        {'code': code, 'date': date},
                        {'$set': doc}, upsert=True
                    )
                )
            except:
                print('发生异常，股票代码：%s，日期：%s' % (code, date), flush=True)
                print(doc, flush=True)

        if len(updateRequests) > 0:
            updateResult = self.db.bulk_write(updateRequests, ordered=False)
            print('抓取股票基本信息，日期：%s, 插入：%4d条，更新：%4d条' %
                  (date, updateResult.upserted_count, updateResult.modified_count), flush=True)

if __name__ == '__main__':
    basicCrawler = BasicCrawler().crawlBasic('2019-06-03', '2019-06-03')